import os  # 导入 os 模块，用于处理文件路径和目录操作
import torch  # 导入 PyTorch 库，用于深度学习相关操作
from torch.utils.data import Dataset  # 导入 PyTorch 的 Dataset 类，用于自定义数据集
from audio_processor import AudioProcessor  # 导入自定义的音频处理工具


class WakeWordDataset(Dataset):  # 定义唤醒词数据集类，继承自 PyTorch 的 Dataset 类
    def __init__(self, data_dir, sr=16000, n_mfcc=13, max_len=100, augment=True):
        """
        初始化方法，接收数据目录、采样率、MFCC 参数、最大长度和是否启用数据增强等参数
        """
        self.sr = sr  # 设置采样率，默认为 16kHz
        self.n_mfcc = n_mfcc  # 设置 MFCC 特征的数量，默认为 13
        self.max_len = max_len  # 设置特征的最大时间步长，默认为 100
        self.augment = augment  # 设置是否启用数据增强
        self.samples = []  # 初始化样本列表，存储每个样本的路径和标签

        # 遍历数据目录
        for label, cls in enumerate(['not_wake', 'wake']):  # 遍历类别（0: not_wake, 1: wake）
            cls_dir = os.path.join(data_dir, cls)  # 构造类别目录路径
            if not os.path.exists(cls_dir):  # 如果目录不存在，跳过
                continue

            for fname in os.listdir(cls_dir):  # 遍历目录中的文件
                path = os.path.join(cls_dir, fname)  # 构造文件完整路径
                self.samples.append((path, label))  # 将文件路径和标签添加到样本列表中

        # 平衡类别
        wake_count = sum(1 for _, label in self.samples if label == 1)  # 统计正样本数量
        not_wake_count = len(self.samples) - wake_count  # 统计负样本数量
        self.weights = [
            1 / wake_count if label == 1 else 1 / not_wake_count  # 计算每个样本的权重
            for _, label in self.samples
        ]

    def __len__(self):
        """返回数据集的样本总数"""
        return len(self.samples)

    def __getitem__(self, idx):
        """获取单个样本，包含错误处理和样本过滤"""
        max_retry = 3  # 最大重试次数，避免因错误样本导致程序中断
        for _ in range(max_retry):  # 循环尝试加载样本，最多重试 max_retry 次
            try:
                path, label = self.samples[idx]  # 获取当前索引对应的样本路径和标签

                # 加载音频
                y, sr = AudioProcessor.load_audio(path, self.sr)  # 使用音频处理器加载音频文件

                # 校验音频长度
                if len(y) < 0.3 * sr:  # 如果音频长度小于 0.3 秒，视为无效样本
                    raise ValueError(f"Audio too short: {len(y) / sr:.2f}s < 0.3s")

                # 数据增强（仅对正样本）
                if self.augment and label == 1:  # 如果启用了数据增强且是正样本
                    y = AudioProcessor.augment_audio(y, sr)  # 对音频进行增强

                # 特征提取
                features = AudioProcessor.extract_features(
                    y, sr, self.n_mfcc, self.max_len
                )  # 提取 MFCC 特征

                # 最终维度校验
                if features.shape[1] != self.max_len:  # 如果特征长度不符合要求
                    raise ValueError(
                        f"Feature length mismatch: {features.shape[1]} vs {self.max_len}"
                    )

                return (
                    torch.FloatTensor(features),  # 返回特征张量 (MFCC, Time)
                    torch.tensor(label, dtype=torch.long),  # 返回标签张量
                    path  # 返回文件路径，便于调试
                )

            except Exception as e:  # 捕获异常
                print(f"Error processing {path} (attempt {_ + 1}/{max_retry}): {str(e)}")  # 打印错误信息
                # 选择下一个样本
                idx = (idx + 1) % len(self)  # 跳过当前样本，尝试下一个样本

        # 重试多次失败后返回空数据
        print(f"Failed after {max_retry} retries, returning zero data")
        return (
            torch.zeros((self.n_mfcc * 3, self.max_len)),  # 返回全零特征张量 (39x100)
            torch.tensor(-1, dtype=torch.long),  # 返回无效标签 -1
            "invalid_sample"  # 返回无效样本标识
        )

    def get_loader(self, batch_size=32, shuffle=True, num_workers=4):
        """获取 DataLoader"""
        sampler = torch.utils.data.WeightedRandomSampler(self.weights, len(self.weights)) if shuffle else None
        # 如果需要打乱数据，使用加权随机采样器；否则不使用
        return torch.utils.data.DataLoader(
            self,
            batch_size=batch_size,  # 设置批量大小
            sampler=sampler,  # 设置采样器
            pin_memory=True,  # 启用内存锁定，加速 GPU 数据传输
            num_workers=num_workers  # 设置多线程工作数
        )
    #这段代码实现了一个自定义的唤醒词数据集类（WakeWordDataset），能够从指定目录加载音频文件，提取特征，并生成可用于训练和评估的 PyTorch 数据加载器。它结合了错误处理、数据增强和加权采样技术，确保数据集的质量和平衡性。